package kim.artemis.rpc.core.registry;

import kim.artemis.rpc.core.common.exception.ArtemisException;
import lombok.extern.slf4j.Slf4j;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

/**
 * Zookeeper 服务注册
 * @author Wq.
 */
@Slf4j
public class ZookeeperServiceRegistry extends ServiceRegistry {

    private CountDownLatch latch = new CountDownLatch(1);

    /**
     * 注册路径
     */
    public static final String ZK_REGISTRY_PATH = "/artemis/registry";

    /**
     * 会话超时间隔
     */
    public static final int ZK_SESSION_TIMEOUT = 5000;

    /**
     * Zk客户端
     */
    private ZooKeeper zkClient;

    /**
     * 刷新发现服务线程
     */
    private Thread refreshDiscoveryThread;

    private volatile boolean refreshStop = false;

    /**
     * 注册服务集合
     */
    private volatile ConcurrentMap<String, TreeSet<String>> registryData;

    /**
     * 发现服务集合
     */
    private volatile ConcurrentMap<String, TreeSet<String>> discoveryData;

    /**
     * 注册路径
     * @param key
     * @return
     */
    public String keyToPath(String key){
        return ZK_REGISTRY_PATH.concat("/").concat(key);
    }

    @Override
    public void start(String address) {
        registryData = new ConcurrentHashMap<>();
        discoveryData = new ConcurrentHashMap<>();
        zkClient = getZkClient(address);
        // 初始化刷新服务发现线程
        refreshDiscoveryThread = new Thread(new Runnable() {
            @Override
            public void run() {
                while (!refreshStop) {
                    try {
                        // 延时1分钟
                        TimeUnit.SECONDS.sleep(60);
                        // 刷新发现服务
                        refreshDiscoveryData(null);
                    } catch (Exception e) {
                        e.printStackTrace();
                        if (!refreshStop) {
                            log.error(">>>>>>> artemis rpc, refresh discovery thread error.");
                        }
                    }
                }
                log.info(">>>>>>> artemis rpc, refresh discovery thread stopd.");
            }
        });
        refreshDiscoveryThread.setName("artemis rpc, ZookeeperServiceRegistry Refresh Discovery Thread");
        refreshDiscoveryThread.setDaemon(true);
        refreshDiscoveryThread.start();
        log.info(">>>>>>> artemis rpc, zookeeper service registry was successfully initialized.");
    }

    @Override
    public void stop () {
        if (zkClient!=null) {
            try {
                zkClient.close();
                zkClient = null;
            } catch (Exception e) {
                log.error(e.getMessage(), e);
            }
        }
        // 停止线程
        if (refreshDiscoveryThread != null) {
            refreshStop = true;
            refreshDiscoveryThread.interrupt();
        }
    }

    /**
     * 连接Zookeeper
     * @param address
     * @return
     */
    private ZooKeeper getZkClient(String address) {
        if (null != zkClient) {
            return zkClient;
        }
        ZooKeeper zooKeeper = null;
        try {
            zooKeeper = new ZooKeeper(address, ZK_SESSION_TIMEOUT, event -> {
                if (event.getState() == Watcher.Event.KeeperState.SyncConnected) {
                    latch.countDown();
                }
            });
            latch.await();
        } catch (IOException | InterruptedException e) {
            log.error(">>>>>>> artemis rpc, Can't get connection for zookeeper,registry address:{}", address, e);
        }
        return zooKeeper;
    }

    /**
     * 使用父路径创建服务持久化路径
     * @param servicePath
     * @param watch
     * @return
     */
    private Stat createServicePathWithParent(String servicePath, boolean watch){
        // valid
        if (null == servicePath || 0 == servicePath.trim().length()) {
            return null;
        }
        try {
            Stat stat = zkClient.exists(servicePath, watch);
            if (stat == null) {
                // 判断是否有上级节点
                if (servicePath.lastIndexOf("/") > 0) {
                    String parentPath = servicePath.substring(0, servicePath.lastIndexOf("/"));
                    Stat parentStat = zkClient.exists(parentPath, watch);
                    if (parentStat == null) {
                        createServicePathWithParent(parentPath, false);
                    }
                }
                zkClient.create(servicePath, new byte[]{}, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
            return zkClient.exists(servicePath, true);
        } catch (Exception e) {
            throw new ArtemisException(e);
        }
    }

    @Override
    public boolean register (Set<String> keys, String value) {
        for (String key : keys) {
            // 本地缓存
            TreeSet<String> values = registryData.get(key);
            if (values == null) {
                values = new TreeSet<>();
                registryData.put(key, values);
            }
            values.add(value);

            if (null != zkClient) {
                try {
                    String servicePath = keyToPath(key);
                    // 创建上级服务路径
                    createServicePathWithParent(servicePath, false);
                    servicePath = servicePath.concat("/address-");
                    zkClient.create(servicePath, value.getBytes("UTF-8"), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
                } catch (Exception e) {
                    throw new ArtemisException(e);
                }
            }
            log.info(">>>>>>> artemis rpc, registry success, keys = {}, value = {}", keys, value);
        }

        return true;
    }

    @Override
    public boolean remove(Set<String> keys, String value) {
        for (String key : keys) {
            TreeSet<String> values = discoveryData.get(key);
            if (values != null) {
                values.remove(value);
            };
            String servicePath = keyToPath(key);
            try {
                Stat stat = zkClient.exists(servicePath, true);
                if (stat != null) {
                    zkClient.delete(servicePath, stat.getVersion());
                } else {
                    log.error(">>>>>>> zookeeper node path not found :{}", servicePath);
                }
            } catch (Exception e) {
                throw new ArtemisException(e);
            }
        }
        log.info(">>>>>>> artemis rpc, remove success, keys = {}, value = {}", keys, value);
        return true;
    }

    @Override
    public boolean remove (String value) {
        if (null == value || 0 == value.trim().length()) {
            return false;
        }
        discoveryData.keySet().forEach(key -> {
            TreeSet<String> values = discoveryData.get(key);
            values.remove(value);
        });
        return true;
    }

    /**
     * 发现服务
     * @param keys
     * @return
     */
    @Override
    public Map<String, TreeSet<String>> discovery(Set<String> keys) {
        if (keys==null || keys.size()==0) {
            return null;
        }
        Map<String, TreeSet<String>> registryDataTmp = new HashMap<>();
        for (String key : keys) {
            TreeSet<String> valueSetTmp = discovery(key);
            if (valueSetTmp != null) {
                registryDataTmp.put(key, valueSetTmp);
            }
        }
        return registryDataTmp;
    }

    /**
     * 发现服务
     * @param key
     * @return
     */
    @Override
    public TreeSet<String> discovery(String key) {
        TreeSet<String> values = discoveryData.get(key);
        if (null == values) {
            refreshDiscoveryData(key);
            values = discoveryData.get(key);
        }
        return values;
    }


    /**
     * 刷新发现服务数据并且更新缓存
     * @param key
     */
    private void refreshDiscoveryData(String key){
        try {
            Set<String> keys = new HashSet<>();
            if (key != null && key.trim().length() > 0) {
                keys.add(key);
            } else {
                if (discoveryData.size() > 0) {
                    keys.addAll(discoveryData.keySet());
                }
            }

            for (String keyItem : keys) {
                String servicePath = keyToPath(keyItem);
                Stat stat = zkClient.exists(servicePath, true);
                if (null == stat) {
                    throw new ArtemisException("artemis rpc, can not find any service node on path: " + servicePath);
                }

                TreeSet<String> addressSet = new TreeSet<>();
                List<String> childrenNode = zkClient.getChildren(servicePath, true);
                for (String node : childrenNode) {
                    String nodePath = servicePath.concat("/").concat(node);

                    byte[] addressData = zkClient.getData(nodePath, false, null);
                    addressSet.add(new String(addressData, "UTF-8"));
                }

                TreeSet<String> existValues = discoveryData.get(keyItem);
                if (null == existValues) {
                    existValues = new TreeSet<>();
                    discoveryData.put(keyItem, existValues);
                }

                if (0 < addressSet.size()) {
                    existValues.clear();
                    existValues.addAll(addressSet);
                }

            }
        } catch (Exception e){
            throw new ArtemisException(e);
        }
    }
}
